module cubed_sphere_process_mod

  use mpi
  use flogger
  use mesh_const_mod
  use cubed_sphere_mesh_mod

  implicit none

  private

  public cubed_sphere_process_type

  type process_neighbor_type
    integer :: id          = MPI_PROC_NULL
    integer :: dom_cart_id = MPI_PROC_NULL
    integer :: orient      = 0
  end type process_neighbor_type

  type cubed_sphere_process_type
    logical :: initialized = .false.
    logical :: mpi_master  = .true.
    integer :: id    = MPI_PROC_NULL
    integer :: comm  = MPI_COMM_WORLD
    integer :: group = MPI_GROUP_NULL
    integer :: np    = 0
    integer :: idom  = 0 ! Domain ID
    integer :: dom_cart_comm  = MPI_COMM_NULL
    integer :: dom_cart_group = MPI_GROUP_NULL
    integer :: dom_cart_id    = MPI_PROC_NULL
    integer :: dom_cart_dims(2)
    integer :: dom_cart_coords(2)
    type(process_neighbor_type), allocatable :: ngb(:)
  contains
    procedure :: init          => cubed_sphere_process_init
    procedure :: decomp_domain => cubed_sphere_process_decomp_domain
    procedure :: clear         => cubed_sphere_process_clear
    procedure :: abort         => cubed_sphere_process_abort
    procedure :: barrier       => cubed_sphere_process_barrier
    procedure :: is_root       => cubed_sphere_process_is_root
    final :: cubed_sphere_process_final
  end type cubed_sphere_process_type

contains

  subroutine cubed_sphere_process_init(this, comm)

    class(cubed_sphere_process_type), intent(inout) :: this
    integer, intent(in), optional :: comm

    integer ierr

    call this%clear()

    if (present(comm)) then
      this%comm = comm
      this%mpi_master = .false.
    else
      call MPI_INIT(ierr)
    end if

    call MPI_COMM_SIZE(this%comm, this%np, ierr)
    call MPI_COMM_RANK(this%comm, this%id, ierr)
    call MPI_COMM_GROUP(this%comm, this%group, ierr)

    if (this%mpi_master) call log_notice('Initialize MPI environment.', pid=this%id)

  end subroutine cubed_sphere_process_init

  subroutine cubed_sphere_process_decomp_domain(this, mesh, proj_type, hdim, vdim, nc, nz, hw, rw, rwv, neq)

    class(cubed_sphere_process_type), intent(inout) :: this
    type(cubed_sphere_mesh_type)    , intent(inout) :: mesh
    integer                         , intent(in   ) :: proj_type
    integer                         , intent(in   ) :: hdim
    integer                         , intent(in   ) :: vdim
    integer                         , intent(in   ) :: nc
    integer                         , intent(in   ) :: nz
    integer                         , intent(in   ) :: hw
    integer                         , intent(in   ) :: rw
    integer                         , intent(in   ) :: rwv
    integer                         , intent(in   ), optional :: neq

    integer dom_np, tmp_comm, tmp_id(2), ierr
    integer i, j, im1, ip1, jm1, jp1, npx, npy
    integer ids, jds, nbx, nby

    if (.not. this%np == 1 .and. mod(this%np, 6) /= 0) then
      stop 'Process size should be divided by 6!'
    end if

    if (.not. allocated(this%ngb)) allocate(this%ngb(8))

    if (this%np == 1) then
      do i = 1, size(this%ngb)
        this%ngb(i)%id = this%id
      end do

      call mesh%init(proj_type=proj_type, hdim=hdim, vdim=vdim, nc=nc, nz=nz, hw=hw, rw=rw, rwv=rwv, neq=neq)
      call mesh%set_ghost_points()
    else
      dom_np = this%np / 6
      this%idom = this%id / dom_np + 1
      call MPI_COMM_SPLIT(this%comm, this%idom, this%id, tmp_comm, ierr)
      call MPI_DIMS_CREATE(dom_np, 2, this%dom_cart_dims, ierr)
      call MPI_CART_CREATE(tmp_comm, 2, this%dom_cart_dims, [.false., .false.], .false., this%dom_cart_comm, ierr)
      call MPI_COMM_GROUP(this%dom_cart_comm, this%dom_cart_group, ierr)
      call MPI_COMM_RANK(this%dom_cart_comm, this%dom_cart_id, ierr)
      call MPI_CART_COORDS(this%dom_cart_comm, this%dom_cart_id, 2, this%dom_cart_coords, ierr)
      call MPI_COMM_FREE(tmp_comm, ierr)

      call MPI_CART_SHIFT(this%dom_cart_comm, 0, 1, this%ngb(left  )%dom_cart_id, this%ngb(right)%dom_cart_id, ierr)
      call MPI_CART_SHIFT(this%dom_cart_comm, 1, 1, this%ngb(bottom)%dom_cart_id, this%ngb(top  )%dom_cart_id, ierr)

      ! Translate Cartesian ID of neighbors to global ID.
      do i = 1, size(this%ngb)
        if (this%ngb(i)%id == MPI_PROC_NULL) then
          call MPI_GROUP_TRANSLATE_RANKS(this%dom_cart_group, 1, [this%ngb(i)%dom_cart_id], this%group, tmp_id, ierr)
          this%ngb(i)%id = tmp_id(1)
        end if
      end do

      ! Set process neighborhood.
      i = this%dom_cart_coords(1)
      j = this%dom_cart_coords(2)
      im1 = i - 1; ip1 = i + 1
      jm1 = j - 1; jp1 = j + 1
      npx = this%dom_cart_dims(1)
      npy = this%dom_cart_dims(2)
      select case (this%idom)
      case (1)
        if (im1 ==  -1) this%ngb(left  )%id = this%id + 3 * dom_np + (npx - 1) * npy
        if (ip1 == npx) this%ngb(right )%id = this%id + npy
        if (jm1 ==  -1) this%ngb(bottom)%id = 5 * dom_np + i * npy + npy - 1
        if (jp1 == npy) this%ngb(top   )%id = 4 * dom_np + i * npy
        if (j >       0 .and. this%ngb(left  )%id /= -1) this%ngb( left_bottom)%id = this%ngb(left  )%id - 1
        if (j < npy - 1 .and. this%ngb(left  )%id /= -1) this%ngb( left_top   )%id = this%ngb(left  )%id + 1
        if (j >       0 .and. this%ngb(right )%id /= -1) this%ngb(right_bottom)%id = this%ngb(right )%id - 1
        if (j < npy - 1 .and. this%ngb(right )%id /= -1) this%ngb(right_top   )%id = this%ngb(right )%id + 1
        if (i >       0 .and. this%ngb(bottom)%id /= -1) this%ngb( left_bottom)%id = this%ngb(bottom)%id - npy
        if (i < npx - 1 .and. this%ngb(bottom)%id /= -1) this%ngb(right_bottom)%id = this%ngb(bottom)%id + npy
        if (i >       0 .and. this%ngb(top   )%id /= -1) this%ngb( left_top   )%id = this%ngb(top   )%id - npy
        if (i < npx - 1 .and. this%ngb(top   )%id /= -1) this%ngb(right_top   )%id = this%ngb(top   )%id + npy
      case (2)
        if (im1 ==  -1) this%ngb(left  )%id = this%id - npy
        if (ip1 == npx) this%ngb(right )%id = this%id + npy
        if (jm1 ==  -1) this%ngb(bottom)%id = 5 * dom_np + (npy - i - 1) + (npx - 1) * npy
        if (jp1 == npy) this%ngb(top   )%id = 4 * dom_np +             i + (npy - 1) * npy
        if (j >       0 .and. this%ngb(left  )%id /= -1) this%ngb( left_bottom)%id = this%ngb(left  )%id - 1
        if (j < npy - 1 .and. this%ngb(left  )%id /= -1) this%ngb( left_top   )%id = this%ngb(left  )%id + 1
        if (j >       0 .and. this%ngb(right )%id /= -1) this%ngb(right_bottom)%id = this%ngb(right )%id - 1
        if (j < npy - 1 .and. this%ngb(right )%id /= -1) this%ngb(right_top   )%id = this%ngb(right )%id + 1
        if (i >       0 .and. this%ngb(bottom)%id /= -1) this%ngb( left_bottom)%id = this%ngb(bottom)%id - merge(npy, -1, j > 0)
        if (i < npx - 1 .and. this%ngb(bottom)%id /= -1) this%ngb(right_bottom)%id = this%ngb(bottom)%id + merge(npy, -1, j > 0)
        if (i >       0 .and. this%ngb(top   )%id /= -1) this%ngb( left_top   )%id = this%ngb(top   )%id - merge(npy,  1, j < npy - 1)
        if (i < npx - 1 .and. this%ngb(top   )%id /= -1) this%ngb(right_top   )%id = this%ngb(top   )%id + merge(npy,  1, j < npy - 1)
      case (3)
        if (im1 ==  -1) this%ngb(left  )%id = this%id - npy
        if (ip1 == npx) this%ngb(right )%id = this%id + npy
        if (jm1 ==  -1) this%ngb(bottom)%id = 5 * dom_np + (npx - i - 1) * npy
        if (jp1 == npy) this%ngb(top   )%id = 4 * dom_np + (npx - i - 1) * npy + j
        if (j >       0 .and. this%ngb(left  )%id /= -1) this%ngb( left_bottom)%id = this%ngb(left  )%id - 1
        if (j < npy - 1 .and. this%ngb(left  )%id /= -1) this%ngb( left_top   )%id = this%ngb(left  )%id + 1
        if (j >       0 .and. this%ngb(right )%id /= -1) this%ngb(right_bottom)%id = this%ngb(right )%id - 1
        if (j < npy - 1 .and. this%ngb(right )%id /= -1) this%ngb(right_top   )%id = this%ngb(right )%id + 1
        if (i >       0 .and. this%ngb(bottom)%id /= -1) this%ngb( left_bottom)%id = this%ngb(bottom)%id - merge(npy, -npy, j > 0)
        if (i < npx - 1 .and. this%ngb(bottom)%id /= -1) this%ngb(right_bottom)%id = this%ngb(bottom)%id + merge(npy, -npy, j > 0)
        if (i >       0 .and. this%ngb(top   )%id /= -1) this%ngb( left_top   )%id = this%ngb(top   )%id - merge(npy, -npy, j < npy - 1)
        if (i < npx - 1 .and. this%ngb(top   )%id /= -1) this%ngb(right_top   )%id = this%ngb(top   )%id + merge(npy, -npy, j < npy - 1)
      case (4)
        if (im1 ==  -1) this%ngb(left  )%id = this%id - npy
        if (ip1 == npx) this%ngb(right )%id = this%id - 3 * dom_np - (npx - 1) * npy
        if (jm1 ==  -1) this%ngb(bottom)%id = 5 * dom_np + i
        if (jp1 == npy) this%ngb(top   )%id = 4 * dom_np + npx - i - 1
        if (j >       0 .and. this%ngb(left  )%id /= -1) this%ngb( left_bottom)%id = this%ngb(left  )%id - 1
        if (j < npy - 1 .and. this%ngb(left  )%id /= -1) this%ngb( left_top   )%id = this%ngb(left  )%id + 1
        if (j >       0 .and. this%ngb(right )%id /= -1) this%ngb(right_bottom)%id = this%ngb(right )%id - 1
        if (j < npy - 1 .and. this%ngb(right )%id /= -1) this%ngb(right_top   )%id = this%ngb(right )%id + 1
        if (i >       0 .and. this%ngb(bottom)%id /= -1) this%ngb( left_bottom)%id = this%ngb(bottom)%id - merge(npy,  1, j > 0)
        if (i < npx - 1 .and. this%ngb(bottom)%id /= -1) this%ngb(right_bottom)%id = this%ngb(bottom)%id + merge(npy,  1, j > 0)
        if (i >       0 .and. this%ngb(top   )%id /= -1) this%ngb( left_top   )%id = this%ngb(top   )%id - merge(npy, -1, j < npy - 1)
        if (i < npx - 1 .and. this%ngb(top   )%id /= -1) this%ngb(right_top   )%id = this%ngb(top   )%id + merge(npy, -1, j < npy - 1)
      case (5)
        if (im1 ==  -1) this%ngb(left  )%id = 3 * dom_np + (npy - j - 1) * npy + npy - 1
        if (ip1 == npx) this%ngb(right )%id = 1 * dom_np + j * npy + npy - 1
        if (jm1 ==  -1) this%ngb(bottom)%id = i * npy + npy - 1
        if (jp1 == npy) this%ngb(top   )%id = 2 * dom_np + (npx - i - 1) * npy + npy - 1
        if (j >       0 .and. this%ngb(left  )%id /= -1) this%ngb( left_bottom)%id = this%ngb(left  )%id + npy
        if (j < npy - 1 .and. this%ngb(left  )%id /= -1) this%ngb( left_top   )%id = this%ngb(left  )%id - npy
        if (j >       0 .and. this%ngb(right )%id /= -1) this%ngb(right_bottom)%id = this%ngb(right )%id - npy
        if (j < npy - 1 .and. this%ngb(right )%id /= -1) this%ngb(right_top   )%id = this%ngb(right )%id + npy
        if (i >       0 .and. this%ngb(bottom)%id /= -1) this%ngb( left_bottom)%id = this%ngb(bottom)%id - npy
        if (i < npx - 1 .and. this%ngb(bottom)%id /= -1) this%ngb(right_bottom)%id = this%ngb(bottom)%id + npy
        if (i >       0 .and. this%ngb(top   )%id /= -1) this%ngb( left_top   )%id = this%ngb(top   )%id - merge(npy, -npy, j < npy - 1)
        if (i < npx - 1 .and. this%ngb(top   )%id /= -1) this%ngb(right_top   )%id = this%ngb(top   )%id + merge(npy, -npy, j < npy - 1)
      case (6)
        if (im1 ==  -1) this%ngb(left  )%id = 3 * dom_np + j * npy
        if (ip1 == npx) this%ngb(right )%id = 1 * dom_np + (npy - j - 1) * npy
        if (jm1 ==  -1) this%ngb(bottom)%id = 2 * dom_np + (npx - i - 1) * npy
        if (jp1 == npy) this%ngb(top   )%id = i * npy
        if (j >       0 .and. this%ngb(left  )%id /= -1) this%ngb( left_bottom)%id = this%ngb(left  )%id - merge(  1,  npy, i > 0)
        if (j < npy - 1 .and. this%ngb(left  )%id /= -1) this%ngb( left_top   )%id = this%ngb(left  )%id + merge(  1,  npy, i > 0)
        if (j >       0 .and. this%ngb(right )%id /= -1) this%ngb(right_bottom)%id = this%ngb(right )%id - merge(  1, -npy, i < npx - 1)
        if (j < npy - 1 .and. this%ngb(right )%id /= -1) this%ngb(right_top   )%id = this%ngb(right )%id + merge(  1, -npy, i < npx - 1)
        if (i >       0 .and. this%ngb(bottom)%id /= -1) this%ngb( left_bottom)%id = this%ngb(bottom)%id - merge(npy, -npy, j > 0)
        if (i < npx - 1 .and. this%ngb(bottom)%id /= -1) this%ngb(right_bottom)%id = this%ngb(bottom)%id + merge(npy, -npy, j > 0)
        if (i >       0 .and. this%ngb(top   )%id /= -1) this%ngb( left_top   )%id = this%ngb(top   )%id - merge(npy,  npy, j < npy - 1)
        if (i < npx - 1 .and. this%ngb(top   )%id /= -1) this%ngb(right_top   )%id = this%ngb(top   )%id + merge(npy,  npy, j < npy - 1)
      end select

      do i = 1, size(this%ngb)
        this%ngb(i)%id = max(-1, this%ngb(i)%id)
      end do

      ! Initialize mesh for the process.
      nbx = nc; nby = nc
      select case (this%idom)
      case (1, 2, 5)
        call round_robin(this%dom_cart_dims(1), this%dom_cart_coords(1),  1, nbx, ids)
        call round_robin(this%dom_cart_dims(2), this%dom_cart_coords(2),  1, nby, jds)
      case (3, 4)
        call round_robin(this%dom_cart_dims(1), this%dom_cart_coords(1), -1, nbx, ids)
        call round_robin(this%dom_cart_dims(2), this%dom_cart_coords(2),  1, nby, jds)
      case (6)
        call round_robin(this%dom_cart_dims(1), this%dom_cart_coords(1),  1, nbx, ids)
        call round_robin(this%dom_cart_dims(2), this%dom_cart_coords(2), -1, nby, jds)
      end select
      call mesh%init(proj_type=proj_type, hdim=hdim, vdim=vdim, nc=nc, nz=nz, hw=hw, rw=rw, rwv=rwv, neq=neq, idom=this%idom, &
                     ids=ids, jds=jds, nbx=nbx, nby=nby, pid=this%id, &
                     cart_dims=this%dom_cart_dims, cart_coords=this%dom_cart_coords, stage=1)
      call mesh%set_ghost_points()
      call mesh%init(proj_type=proj_type, hdim=hdim, vdim=vdim, nc=nc, nz=nz, hw=hw, rw=rw, rwv=rwv, neq=neq, idom=this%idom, &
                     ids=ids, jds=jds, nbx=nbx, nby=nby, pid=this%id, &
                     cart_dims=this%dom_cart_dims, cart_coords=this%dom_cart_coords, stage=2)
      call mesh%set_ghost_points()
    end if

  end subroutine cubed_sphere_process_decomp_domain

  pure logical function cubed_sphere_process_is_root(this) result(res)

    class(cubed_sphere_process_type), intent(in) :: this

    res = this%id == 0

  end function cubed_sphere_process_is_root

  subroutine cubed_sphere_process_abort(this)

    class(cubed_sphere_process_type), intent(inout) :: this

    integer ierr

    call MPI_BARRIER(this%comm, ierr)
    if (this%dom_cart_comm /= MPI_COMM_NULL) call MPI_COMM_FREE(this%dom_cart_comm, ierr)
    call MPI_FINALIZE(ierr)

  end subroutine cubed_sphere_process_abort

  subroutine cubed_sphere_process_barrier(this)

    class(cubed_sphere_process_type), intent(in) :: this

    integer ierr

    call MPI_BARRIER(this%comm, ierr)

  end subroutine cubed_sphere_process_barrier

  subroutine cubed_sphere_process_clear(this)

    class(cubed_sphere_process_type), intent(inout) :: this

    integer ierr

    this%comm = MPI_COMM_WORLD
    this%np   = 0
    if (allocated(this%ngb)) deallocate(this%ngb)

  end subroutine cubed_sphere_process_clear

  subroutine cubed_sphere_process_final(this)

    type(cubed_sphere_process_type), intent(inout) :: this

    integer ierr

    if (this%mpi_master) then
      call this%barrier()
      call MPI_FINALIZE(ierr)
      call log_notice('Finalize MPI environment.', pid=this%id)
    end if

  end subroutine cubed_sphere_process_final

  subroutine round_robin(dim, coord, order, num, ibeg, iend)

    integer, intent(in) :: dim
    integer, intent(in) :: coord
    integer, intent(in) :: order
    integer, intent(inout) :: num
    integer, intent(out) :: ibeg ! Start from 1.
    integer, intent(out), optional :: iend ! Start from 1.

    integer res_num, tmp_num, i

    res_num = mod(num, dim)
    ibeg = 1
    select case (order)
    case (1)
      do i = 0, coord - 1
        if (res_num /= 0 .and. i < res_num) then
          tmp_num = num / dim + 1
        else
          tmp_num = num / dim
        end if
        ibeg = ibeg + tmp_num
      end do
      if (res_num /= 0 .and. coord < res_num) then
        num = num / dim + 1
      else
        num = num / dim
      end if
    case (-1)
      do i = 0, coord - 1
        if (res_num /= 0 .and. i >= dim - res_num) then
          tmp_num = num / dim + 1
        else
          tmp_num = num / dim
        end if
        ibeg = ibeg + tmp_num
      end do
      if (res_num /= 0 .and. coord >= dim - res_num) then
        num = num / dim + 1
      else
        num = num / dim
      end if
    end select
    if (present(iend)) iend = ibeg + num - 1

  end subroutine round_robin

end module cubed_sphere_process_mod
